from collections import defaultdict, deque
from typing import Dict, Generator, Iterable, List, Optional, Set, Tuple, TypeAlias, Union

import networkx as nx


def all_depended_slices(slices: Iterable["CodeSlice"]) -> Set["CodeSlice"]:
    slices = list(slices)
    worklist = deque(slices)
    _slices = set()
    while worklist:
        s = worklist.popleft()
        if s in _slices:
            continue
        _slices.add(s)
        for d in s.depends:
            worklist.append(d)
    return _slices.difference(slices)


CodeSliceID: TypeAlias = str


class CodeSlice(object):
    """
    A code slice is the smallest unit of code that we can provide. A slice contains at least one top-level declaration of entity.
    Here "top-level" means the entities declared inside another entity are not to individual slices, but contained in their outer entities' slices
    Entities include global variables, functions(and member functions), user-types, macros, but not namespaces and linkage specifiers.
    So the stuffs inside the namespaces or linkage specifiers are considered as top-level.

    Attributes:
        loc (str): source location and the id of the code slice, in the format of "file:start_line-end_line"
        file (str): the file path of the code slice
        start_line (int): the start line of the code slice
        end_line (int): the end line of the code slice
        decls (List[Tuple[str, str]]): a list of (declaration type, name) in the code slice. Declarations include:
            - Definitions of functions, excluding the compiler generated constructors, operators.
            - Definitions of global variables. Supposed there is no global variable generated by compiler.
            - Declarations of user-types. Notice that a user-type can be "defined" multiple times in different files as long as it appears only once in the source file.
            This will create multiple slices and they are depended on at the same time. Let's wish the source code not so ugly!
            TODO: make multiple declarations of the same user-type as one slice. Qualified names are the identities of types!
            - Macros
            - Not template initializations.
        depends (Set[CodeSlice]): a set of code slices that this code slice depends on. "Depend" means a code slice uses something declared in other code slices.
            - The definition of declaration is depended on if there is.
            - Field access is also a dependency.
            - No self-dependency.
        depended (Set[CodeSlice]): a set of code slices that depends on this code.
        extra: some other info about this slice.
        level: negative number, the smaller the lower, topo order of this slice among all slices.
        text: the text of the code slice
    """

    pool: Optional["CodeSlicePool"]
    file: str
    start_line: int
    end_line: int
    decls: List[Tuple[str, str]]
    text: str
    w_depends: Dict[CodeSliceID, int]
    w_depended: Dict[CodeSliceID, int]
    level: int
    extra: Dict["str", List]
    size: int

    def __init__(self, pool: Optional["CodeSlicePool"], file: str, start_line: int, end_line: int):
        self.pool = pool
        self.file: str = file
        self.start_line: int = start_line
        self.end_line: int = end_line
        self.decls: List[Tuple[str, str]] = []
        self.text: str = ""
        self.w_depends: Dict[CodeSliceID, int] = defaultdict(int)
        self.w_depended: Dict[CodeSliceID, int] = defaultdict(int)
        self.extra = defaultdict(list)
        self.level = 0
        self.size = 0

        # FIXME: 一个已知的bug，也许是bug，file是绝对路径，当pkl存档在机器之间迁移，源代码文件换路径的时候，就会出错

    @property
    def id(self) -> CodeSliceID:
        if self.pool is None:
            namespace = "CodeSlice"
        else:
            namespace = self.pool.namespace
        return f"{namespace}:{self.file}:{self.start_line}-{self.end_line}"

    @property
    def loc(self) -> str:
        return f"{self.file}:{self.start_line}-{self.end_line}"

    def cover(self, other: "CodeSlice") -> bool:
        if not isinstance(other, CodeSlice):
            return NotImplemented("CodeSlice-cover: other must be a CodeSlice")
        return self.start_line <= other.start_line and other.end_line <= self.end_line

    def __hash__(self):
        if hasattr(self, "id"):
            return hash(self.id)
        else:
            return super(CodeSlice, self).__hash__()

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, CodeSlice):
            return NotImplemented
        return self.loc == other.loc

    def __lt__(self, other: object) -> bool:
        if not isinstance(other, CodeSlice):
            return NotImplemented
        return self.loc < other.loc

    def __repr__(self) -> str:
        return (
            f"CodeSlice(file='{self.file}', start_line={self.start_line}, " f"end_line={self.end_line}, " f"decls={self.decls})"
        )

    @property
    def depends(self) -> Set["CodeSlice"]:
        if self.pool is None:
            raise ValueError("Cannot add CodeSlice as a dependency from None pool.")
        return {self.pool.get_slice_by_id(id) for id in self.w_depends.keys()}

    @property
    def depended(self) -> Set["CodeSlice"]:
        if self.pool is None:
            raise ValueError("Cannot add CodeSlice as a dependency from None pool.")
        return {self.pool.get_slice_by_id(id) for id in self.w_depended.keys()}

    def isMain(self):
        return any(decl[1] == "main" for decl in self.decls)

    def isTestFunc(self):
        return any(decl[1].startswith("test_") for decl in self.decls)

    def isTypedef(self):
        return any(decl[0] == "CTypedefType" for decl in self.decls)

    def depends_on(self, dependency: Union["CodeSlice", CodeSliceID]):
        """
        This method add dependency of this code slice. Do nothing if the dependency is itself.
        """
        if self.pool is None:
            raise ValueError("Cannot depends_on CodeSlice from None pool.")
        if isinstance(dependency, CodeSliceID):
            dependency = self.pool.get_slice_by_id(dependency)

        if dependency == self:
            return

        self.w_depends[dependency.id] += 1

        dependency.w_depended[self.id] += 1

    def wont_depend_on(self, dependency: Union["CodeSlice", CodeSliceID]):
        """
        This method remove dependency of this code slice. Do nothing if the dependency doesn't exist.
        """
        if self.pool is None:
            raise ValueError("Cannot wont_depend_on CodeSlice from None pool.")
        if isinstance(dependency, CodeSliceID):
            dependency = self.pool.get_slice_by_id(dependency)

        if dependency not in self.depends:
            return

        del self.w_depends[dependency.id]
        del dependency.w_depended[self.id]


class CodeSlicePool:
    namespace: str
    code_slices: Dict[str, List[CodeSlice]]

    def __init__(self, namespace: str = "CodeSlice") -> None:
        self.namespace = namespace
        self.code_slices: Dict[str, List[CodeSlice]] = defaultdict(list)

    def get_slice_by_id(self, id: CodeSliceID) -> CodeSlice:
        ret = None
        for file_name, slices in self.code_slices.items():
            for slice in slices:
                if slice.id == id:
                    ret = slice
                    break
            if ret is not None:
                break
        if ret is None:
            raise KeyError(f"Slice {id} not found")
        return ret

    def all_code_slices(self) -> Generator[CodeSlice, None, None]:
        for slices in self.code_slices.values():
            yield from slices

    def add_code_slice(self, code: CodeSlice) -> bool:
        to_remove = []
        for code_slice in self.code_slices[code.file]:
            if code_slice.cover(code):
                return False
            if code.cover(code_slice):
                to_remove.append(code_slice)
        for code_slice in to_remove:
            self.code_slices[code.file].remove(code_slice)
        self.code_slices[code.file].append(code)
        code.pool = self
        return True

    def find_code_slice(self, file: str, start_line: int, end_line: int) -> Optional[CodeSlice]:
        for code in self.code_slices[file]:
            if code.start_line <= start_line <= end_line <= code.end_line:
                return code
        return None

    def seal(self, size_calc=len) -> None:
        """
        Call this when all code slices are all collected, this will remove the code slices without declarations. And fill the text of the code slices.
        """

        def reserved(slice: CodeSlice):
            if not slice.decls:
                return False
            if slice.isMain():
                return False
            if any(decl == ("GlobalVariable", "tests") for decl in slice.decls):
                return False
            if any(decl == ("TopLevelFunction", "run_tests") for decl in slice.decls):
                return False
            if any(decl == ("TopLevelFunction", "run_test") for decl in slice.decls):
                return False
            if any(decl == ("CTypedefType", "UnitTestFunction") for decl in slice.decls):
                return False
            if all(decl[0] == "Macro" for decl in slice.decls) and len(slice.depended) == 0:
                return False
            return True

        code_slices = self.code_slices
        new_code_slices = {}
        for file, slices in code_slices.items():
            new_code_slices[file] = []
            for slice in slices:
                if not reserved(slice):
                    for dep in slice.depends:
                        slice.wont_depend_on(dep)
                    for dep in slice.depended:
                        dep.wont_depend_on(slice)
                    continue
                new_code_slices[file].append(slice)
        self.code_slices = new_code_slices

        # set level
        g = nx.DiGraph()
        g.add_edges_from((s, d) for s in self.all_code_slices() for d in s.depends)
        comp = list(nx.strongly_connected_components(g))
        comp_of_slice = {s: i for i, c in enumerate(comp) for s in c}
        dag = nx.DiGraph()
        dag.add_edges_from(
            (s2, d2) for s in self.all_code_slices() for d in s.depends if (s2 := comp_of_slice[s]) != (d2 := comp_of_slice[d])
        )
        for i, comps in enumerate(nx.topological_generations(dag)):
            for c in comps:
                for s in comp[c]:
                    s.level = -1 - i
        for file, slices in self.code_slices.items():
            with open(file, "r") as f:
                lines = f.readlines()
            for slice in slices:
                if not slice.text:
                    slice.text = "".join(lines[slice.start_line - 1 : slice.end_line])
                slice.size = size_calc(slice.text)
